Part of NVIDIA-Merlin/Merlin#733.

Goals ⚽

There is a package called distributed-embeddings, a library for building large embedding based (e.g. recommender) models in Tensorflow. It's an alternative approach to SOK.

This PR introduces DistributedEmbedding for multi-GPU embedding table support.

Implementation Details 🚧

• distributed-embeddings by default will round-robin the entire embedding tables across the GPUs, e.g., the first embedding table on GPU 1, the second one on GPU 2, etc.
• In theory the tables can be sharded by using column_slice but this has not been tested thoroughly from Models side.
• Most of the logic is for inferring the embedding size from the schema using the cardinality in int_domain (similarly to the existing EmbeddingTable), determining shapes, and translating a dictionary input into an ordered list input (because distributed-embeddings doesn't support dictionaries yet).
• From the user perspective, they can replace mm.Embeddings with mm.DistributedEmbeddings in their models when they wish to use multi-GPU embedding tables. (See the unit test for DLRM.)
• Depends on upstream fix: Use tf.shape for graph mode support distributed-embeddings#6
• Added a Github actions for running unit tests that depend on horovod. distributed-embeddings is for now installed via a script that clones the github repo and installs from source, because there is no pypi package.

Testing Details 🔍

Unit tests: tests/unit/tf/horovod/test_embedding.py
Performance tests: TBD